import win32gui
import win32ui
import win32con
from ctypes import windll
import numpy as np
from PIL import Image

def capture_mumu_window():
    # 1. 获取窗口句柄
    hwnd = 722998
    if not hwnd:
        raise ValueError("未找到MuMu窗口")

    # 2. 获取窗口尺寸
    left, top, right, bottom = win32gui.GetClientRect(hwnd)
    width = right - left
    height = bottom - top

    # 3. 创建设备上下文
    hwnd_dc = win32gui.GetWindowDC(hwnd)
    mfc_dc = win32ui.CreateDCFromHandle(hwnd_dc)
    save_dc = mfc_dc.CreateCompatibleDC()

    # 4. 创建位图
    save_bitmap = win32ui.CreateBitmap()
    save_bitmap.CreateCompatibleBitmap(mfc_dc, width, height)
    save_dc.SelectObject(save_bitmap)

    # 5. 截图（优先PrintWindow）
    result = windll.user32.PrintWindow(hwnd, save_dc.GetSafeHdc(), 0)
    if not result:
        # 若PrintWindow失败，尝试BitBlt
        save_dc.BitBlt((0, 0), (width, height), mfc_dc, (0, 0), win32con.SRCCOPY)

    # 6. 转换为PIL图像
    bmp_info = save_bitmap.GetInfo()
    bmp_bits = save_bitmap.GetBitmapBits(True)
    frame = np.frombuffer(bmp_bits, dtype=np.uint8).reshape((height, width, 4))
    image = Image.fromarray(frame).convert("RGB")
    image.save("mumu_screenshot.png")

    # 7. 清理资源
    win32gui.DeleteObject(save_bitmap.GetHandle())
    save_dc.DeleteDC()
    mfc_dc.DeleteDC()
    win32gui.ReleaseDC(hwnd, hwnd_dc)

    return image

# 执行截图
capture_mumu_window()
